import time

import pymysql
import pymysql.cursors

from static.data.process.datapreprocess import process
import cv2
from anoInversion import llama
import torch
from PIL import Image

from pytorch_lightning import Trainer
from torch.utils.data import DataLoader

from anoDetection.anomalib.config import get_configurable_parameters
from anoDetection.anomalib.data.inference import InferenceDataset
from anoDetection.anomalib.data.utils import InputNormalizationMethod, get_transforms
from anoDetection.anomalib.models import get_model
from anoDetection.anomalib.utils.callbacks import get_callbacks
from anoGeneration.generation_inference import generate_image

from Inversion.inversion import inversion

device = "cuda" if torch.cuda.is_available() else "cpu"

db_config = {
    "host": "localhost",
    "user": "root",
    "password": "123456",
    "db": "digital_twin_local",
}


def connect_to_database():
    connection = pymysql.connect(cursorclass=pymysql.cursors.DictCursor, **db_config)
    return connection


def select_unchecked_image_data():
    connection = connect_to_database()
    try:
        with connection.cursor() as cursor:
            sql = "select timetamp, folder_name, img_path, format from tb_image where checked = 0"
            print(sql)
            cursor.execute(sql)
            result = cursor.fetchall()
            connection.close()
            return result
    finally:
        # 注意：这里不关闭连接，保持连接的唯一性
        pass


def select_unchecked_text_data():
    connection = connect_to_database()
    try:
        with connection.cursor() as cursor:
            sql = "select temperature, humidity, air_pressure, wind_speed,wind_direction, total_cloud, flew_cloud, " \
                  "cloud_form, visibility, precipitation, timetamp, bv, swd, swu, lwd, lwu, ts, sd, pt " \
                  "from tb_text ORDER BY timetamp DESC LIMIT 1"
            print(sql)
            cursor.execute(sql)
            result = cursor.fetchall()
            connection.close()
            return result
    finally:
        # 注意：这里不关闭连接，保持连接的唯一性
        pass


def update_unchecked_image_data(timetamp):
    connection = connect_to_database()
    try:
        with connection.cursor() as cursor:
            sql = "update tb_image set checked = 1  where timetamp = '%s'" % (timetamp)
            print(sql)
            cursor.execute(sql)
            connection.commit()
            connection.close()
    finally:
        # 注意：这里不关闭连接，保持连接的唯一性
        pass


def update_unchecked_text_data(timetamp):
    connection = connect_to_database()
    try:
        with connection.cursor() as cursor:
            sql = "update tb_text set checked = 1  where timetamp = '%s'" % (timetamp)
            print(sql)
            cursor.execute(sql)
            connection.commit()
            connection.close()
    finally:
        # 注意：这里不关闭连接，保持连接的唯一性
        pass


def update_result_data(key, value, detect_image):
    connection = connect_to_database()
    try:
        with connection.cursor() as cursor:
            sql = "update tb_result set %s = '%s' where detect_image = '%s'" % (key, value, detect_image)
            print(sql)
            cursor.execute(sql)
            connection.commit()
            connection.close()
    finally:
        # 注意：这里不关闭连接，保持连接的唯一性
        pass


def create_prompt(text_result, i):
    # id = text_result[i]['id']
    temperature = text_result[i]['temperature']
    humidity = text_result[i]['humidity']
    air_pressure = text_result[i]['air_pressure']
    wind_speed = text_result[i]['wind_speed']
    wind_direction = text_result[i]['wind_direction']
    total_cloud = text_result[i]['total_cloud']
    flew_cloud = text_result[i]['flew_cloud']
    cloud_form = text_result[i]['cloud_form']
    visibility = text_result[i]['visibility']
    precipitation = text_result[i]['precipitation']
    timetamp = text_result[i]['timetamp']
    bv = text_result[i]['bv']
    swd = text_result[i]['swd']
    swu = text_result[i]['swu']
    lwd = text_result[i]['lwd']
    lwu = text_result[i]['lwu']
    ts = text_result[i]['ts']
    sd = text_result[i]['sd']
    pt = text_result[i]['pt']
    # checked = text_result[i]['checked']
    prompt = "temperature: " + str(temperature) + ", humidity: " + str(humidity) + ", air_pressure: " \
             + str(air_pressure) + ", wind_speed: " + str(wind_speed) + ", wind_direction: " \
             + str(wind_direction) + ", total_cloud: " + str(total_cloud) + ", flew_cloud: " + str(flew_cloud) \
             + ", cloud_form: n" + str(cloud_form) + ", visibility: " + str(visibility) + ", precipitation: " \
             + str(precipitation) + ", time: " + str(timetamp) + ", bv: " + str(bv) + ", swd: " + str(swd) \
             + ", swu: " + str(swu) + ", lwd: " + str(lwd) + ", lwu: " + str(lwu) + ", ts: " + str(ts) \
             + ", sd: " + str(sd) + ", pt: " + str(pt)
    prompt.replace('None', 'null')
    return prompt


def load():
    config = get_configurable_parameters(config_path='./anoDetection/anomalib/models/patchcore/config.yaml')
    config.trainer.resume_from_checkpoint = './anoDetection/results/patchcore/mdata/spectrogram8/run/weights/lightning/model-v3.ckpt'
    config.visualization.save_images = True
    config.visualization.image_save_path = './static/detect_results/mdata'

    # create model and trainer
    model = get_model(config)
    callbacks = get_callbacks(config)
    trainer = Trainer(callbacks=callbacks, **config.trainer)

    # get the transforms
    transform_config = config.dataset.transform_config.eval if "transform_config" in config.dataset.keys() else None
    image_size = (config.dataset.image_size[0], config.dataset.image_size[1])
    center_crop = config.dataset.get("center_crop")
    if center_crop is not None:
        center_crop = tuple(center_crop)
    normalization = InputNormalizationMethod(config.dataset.normalization)
    transform = get_transforms(
        config=transform_config, image_size=image_size, center_crop=center_crop, normalization=normalization
    )
    return config, transform, trainer, model


def infer(img_input, config, transform, trainer, model):
    dataset = InferenceDataset(
        img_input, image_size=tuple(config.dataset.image_size), transform=transform  # type: ignore
    )
    dataloader = DataLoader(dataset)

    # generate predictions
    trainer.predict(model=model, dataloaders=[dataloader])

def create_report():
    connection = connect_to_database()

    try:
        with connection.cursor() as cursor:
            sql = "select * from tb_result ORDER BY timetamp DESC LIMIT 1"
            print(sql)
            cursor.execute(sql)
            result_list = cursor.fetchall()
            connection.close()
    finally:
        # 注意：这里不关闭连接，保持连接的唯一性
        pass

    result = result_list[0]
    print(result)
    print(result_list)
    if int(result.get('detect_label')) == 0:
        detect_label_report = '图像受环境干扰因素较小，环境干扰指数为%s。' % (result.get('detect_score'))
        detect_report = detect_label_report
    else:
        detect_label_report = '图像受环境干扰因素较大，环境干扰指数为%s。' % (result.get('detect_score'))
        if float(result.get('detect_score')) > 60:
            detect_score_report = '。此时图像受环境干扰情况严重，图像拍摄效果较差，建议弃用。'
            detect_report = detect_label_report + detect_score_report
        elif float(result.get('detect_score')) < 60:
            detect_score_report = '。此时图像受环境干扰情况较重，图像拍摄效果受到影响，检测效果与实际结果存在一定误差，建议谨慎使用。'
            detect_report = detect_label_report + detect_score_report

    if int(result.get('inversion_score')) == -1:
        inversion_score_report = '此时拍摄图像强度较弱，不适用于使用反演算法计算反演温度，可能存在以下问题的可能：' \
                                 '1.此时图像与暗噪声图像不适配；2.提取背景杂散光出现问题。建议调整人工调整算法参数计算反演温度，' \
                                 '或重新拍摄暗噪声图像与现图像匹配。'
    else:
        inversion_score_report = '此时反演算法计算的反演温度为%s。' % (result.get('inversion_score'))

    report = '此结果生成时间为%s，该时刻' % (result.get('timetamp')) + detect_report + inversion_score_report + '该时刻' \
                                                                                                               '的图像受环境干扰结果图，转动温度光谱预处理后图像，预测图像分别保存在%s，%s，%s。' \
             % (result.get('detect_image'), result.get('spectrogram8'), result.get('generation_image'))

    return report



if __name__ == '__main__':
    config, transform, trainer, detection_model = load()
    print("load detection_model successfully")
    llama_dir = "anoInversion/"
    inversion_model, preprocess = llama.load("anoInversion/result/checkpoint-0.pth", llama_dir, llama_type="7B",
                                             device=device)
    inversion_model.eval()
    print("load inversion_model successfully")

    while True:
        print("轮询开始")
        # run(config, transform, trainer, detection_model, inversion_model, preprocess)
        img_result = select_unchecked_image_data()
        text_result = select_unchecked_text_data()


        for i in range(len(img_result)):
            img_timetamp = img_result[i]['timetamp']
            folder_name = img_result[i]['folder_name']
            img_path = img_result[i]['img_path']
            img_format = img_result[i]['format']

            # img_input = 'data/spectrogram/%s/%s.%s' % (folder_name, img_path, img_format)

            # 异常检测
            img8_input, img16_input = process(folder_name, img_path, img_format)
            infer(img8_input, config, transform, trainer, detection_model)

            # 异常检测结果写入result
            detect_image = "static/detect_results/mdata/spectrogram8/%s/%s.png" % (folder_name, img_path)
            update_result_data('timetamp', img_timetamp, detect_image)
            update_unchecked_image_data(img_timetamp)

            # 预处理后8位图像写入result
            spectrogram8 = 'static/data/spectrogram8/%s/%s.png' % (folder_name, img_path)
            update_result_data('spectrogram8', spectrogram8, detect_image)

            # 16位图像送入反演算法中计算，并将结果写入result
            try:
                inversion_score = inversion(img16_input)
                print(inversion_score)
            except:
                inversion_score = '-1'
            update_result_data('inversion_score', inversion_score, detect_image)

            # 生成prompt，与8位图像一同送入反演模型，计算反演温度，结果写入result
            prompt = create_prompt(text_result, 0)
            print(prompt)
            print(img8_input)
            prompt = llama.format_prompt(prompt)
            img = Image.fromarray(cv2.imread(img8_input))
            img = preprocess(img).unsqueeze(0).to(device)
            anoinv_score = inversion_model.generate(img, [prompt])[0]
            print(anoinv_score)
            update_result_data('anoinv_score', anoinv_score, detect_image)
            update_unchecked_text_data(text_result[0]['timetamp'])

            # 读取数据库中最近的10张图像，预测生成下一时刻的图像，结果写入result
            generation_image, sum_image, generation_BG_image = generate_image(img_timetamp)
            update_result_data('generation_image', generation_image, detect_image)
            update_result_data('sum_image', sum_image, detect_image)
            update_result_data('generation_BG_image', generation_BG_image, detect_image)

            # 根据已有结果给出结果报告
            report = create_report()
            update_result_data('report', report, detect_image)

        time.sleep(1)
